#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from typing import Dict, List, Set
from dataclasses import dataclass
from enum import Enum

class VPNProtocol(Enum):
    OPENVPN = "OpenVPN"
    IPSEC = "IPSec/IKEv2"
    WIREGUARD = "WireGuard"
    SSTP = "SSTP"
    L2TP = "L2TP"
    PPTP = "PPTP"
    CISCO_VPN = "Cisco VPN"
    ANYCONNECT = "AnyConnect"
    SOFETHER = "SoftEther"
    SHADOWSOCKS = "Shadowsocks"
    V2RAY = "V2Ray"
    TROJAN = "Trojan"
    TOR = "Tor"
    SOCKS5 = "Socks5"
    SSH_TUNNEL = "SSH Tunnel"
    FORTIGATE = "FortiGate SSL VPN"
    PALO_ALTO = "Palo Alto GlobalProtect"
    PULSE_SECURE = "Pulse Secure"
    NORDLYNX = "NordLynx"
    TAILSCALE = "Tailscale"

@dataclass
class ProtocolFeature:
    """VPN协议特征定义"""
    name: VPNProtocol
    ports: Set[int]
    packet_patterns: List[Dict]
    tls_fingerprints: List[Dict]
    flow_characteristics: Dict
    detection_difficulty: str
    evasion_techniques: List[str]

# VPN协议特征知识库
VPN_KNOWLEDGE_BASE = {
    VPNProtocol.OPENVPN: ProtocolFeature(
        name=VPNProtocol.OPENVPN,
        ports={1194, 443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "UDP", "length": 148},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256", "TLS_AES_256_GCM_SHA384"]},
            {"extensions": ["server_name", "status_request", "supported_groups"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 0.7, "UDP": 0.3},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="中等",
        evasion_techniques=["UDP模式混淆", "动态端口切换", "TLS伪装"]
    ),
    
    VPNProtocol.IPSEC: ProtocolFeature(
        name=VPNProtocol.IPSEC,
        ports={500, 4500},
        packet_patterns=[
            {"layer": "ISAKMP", "exchange_type": 1},
            {"layer": "ESP", "protocol": 50},
            {"layer": "UDP", "length": 1024}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"UDP": 0.8, "ESP": 0.2},
            "packet_interval": {"min": 0.05, "max": 0.5}
        },
        detection_difficulty="高",
        evasion_techniques=["NAT穿越流量", "ESP封装", "动态端口"]
    ),
    
    VPNProtocol.WIREGUARD: ProtocolFeature(
        name=VPNProtocol.WIREGUARD,
        ports={51820},
        packet_patterns=[
            {"layer": "UDP", "length": 148},
            {"layer": "UDP", "length": 92}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 92,
            "max_packet_size": 148,
            "protocol_ratio": {"UDP": 1.0},
            "packet_interval": {"min": 0.1, "max": 0.3}
        },
        detection_difficulty="高",
        evasion_techniques=["无协议头特征", "固定包大小", "UDP封装"]
    ),
    
    VPNProtocol.SSTP: ProtocolFeature(
        name=VPNProtocol.SSTP,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="中等",
        evasion_techniques=["HTTPS隧道", "SSL协商", "动态端口"]
    ),
    
    VPNProtocol.L2TP: ProtocolFeature(
        name=VPNProtocol.L2TP,
        ports={1701},
        packet_patterns=[
            {"layer": "UDP", "length": 1701},
            {"layer": "IPSec", "protocol": 50}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"UDP": 0.6, "ESP": 0.4},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="高",
        evasion_techniques=["双协议栈识别", "NAT穿越", "动态端口"]
    ),
    
    VPNProtocol.PPTP: ProtocolFeature(
        name=VPNProtocol.PPTP,
        ports={1723},
        packet_patterns=[
            {"layer": "TCP", "flags": "S", "options": ["MSS"]},
            {"layer": "GRE", "protocol": 47}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 0.3, "GRE": 0.7},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="低",
        evasion_techniques=["企业合法使用", "GRE封装", "动态端口"]
    ),
    
    VPNProtocol.CISCO_VPN: ProtocolFeature(
        name=VPNProtocol.CISCO_VPN,
        ports={443},
        packet_patterns=[
            {"layer": "DTLS", "handshake_type": "ClientHello"},
            {"layer": "UDP", "length": 1024}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"UDP": 0.8, "TCP": 0.2},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="中等",
        evasion_techniques=["企业合法流量", "DTLS握手", "动态端口"]
    ),
    
    VPNProtocol.ANYCONNECT: ProtocolFeature(
        name=VPNProtocol.ANYCONNECT,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request", "signature_algorithms"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="高",
        evasion_techniques=["独特证书OID", "CSTP协议头", "加密深度混淆"]
    ),
    
    VPNProtocol.SOFETHER: ProtocolFeature(
        name=VPNProtocol.SOFETHER,
        ports={443, 992, 1194, 5555},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 0.8, "UDP": 0.2},
            "packet_interval": {"min": 0.1, "max": 0.3}
        },
        detection_difficulty="高",
        evasion_techniques=["伪装HTTPS", "特定心跳包间隔", "动态端口切换"]
    ),
    
    VPNProtocol.SHADOWSOCKS: ProtocolFeature(
        name=VPNProtocol.SHADOWSOCKS,
        ports={8388},
        packet_patterns=[
            {"layer": "TCP", "flags": "S", "options": ["MSS"]},
            {"layer": "TCP", "length": 16}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 16,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.05, "max": 0.2}
        },
        detection_difficulty="高",
        evasion_techniques=["固定头结构", "高熵载荷", "AEAD加密模式"]
    ),
    
    VPNProtocol.V2RAY: ProtocolFeature(
        name=VPNProtocol.V2RAY,
        ports={10086},
        packet_patterns=[
            {"layer": "WebSocket", "opcode": 1},
            {"layer": "TCP", "flags": "S", "options": ["MSS"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="高",
        evasion_techniques=["VMess协议ID", "伪装WebSocket", "多重传输层封装"]
    ),
    
    VPNProtocol.TROJAN: ProtocolFeature(
        name=VPNProtocol.TROJAN,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="高",
        evasion_techniques=["TLS+特洛伊协议头", "全流量TLS加密", "动态端口"]
    ),
    
    VPNProtocol.TOR: ProtocolFeature(
        name=VPNProtocol.TOR,
        ports={9001},
        packet_patterns=[
            {"layer": "TCP", "flags": "S", "options": ["MSS"]},
            {"layer": "TCP", "length": 512}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 512,
            "max_packet_size": 512,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 0.3}
        },
        detection_difficulty="高",
        evasion_techniques=["中继节点指纹", "Cell包结构", "流量洋葱路由"]
    ),
    
    VPNProtocol.SOCKS5: ProtocolFeature(
        name=VPNProtocol.SOCKS5,
        ports={1080},
        packet_patterns=[
            {"layer": "TCP", "flags": "S", "options": ["MSS"]},
            {"layer": "TCP", "length": 3}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 3,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="低",
        evasion_techniques=["认证协商模式", "非加密模式识别", "动态端口"]
    ),
    
    VPNProtocol.SSH_TUNNEL: ProtocolFeature(
        name=VPNProtocol.SSH_TUNNEL,
        ports={22},
        packet_patterns=[
            {"layer": "TCP", "flags": "S", "options": ["MSS"]},
            {"layer": "SSH", "version": "SSH-2.0"}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="中等",
        evasion_techniques=["SSH版本协商", "端口转发模式", "与运维流量混淆"]
    ),
    
    VPNProtocol.FORTIGATE: ProtocolFeature(
        name=VPNProtocol.FORTIGATE,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="中等",
        evasion_techniques=["特定Cookie值", "FGT标识", "企业混合流量"]
    ),
    
    VPNProtocol.PALO_ALTO: ProtocolFeature(
        name=VPNProtocol.PALO_ALTO,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="中等",
        evasion_techniques=["GP头标识", "移动端特征", "动态端口"]
    ),
    
    VPNProtocol.PULSE_SECURE: ProtocolFeature(
        name=VPNProtocol.PULSE_SECURE,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 1.0},
            "packet_interval": {"min": 0.1, "max": 1.0}
        },
        detection_difficulty="中等",
        evasion_techniques=["DSID Cookie特征", "Java签名", "会话保持机制"]
    ),
    
    VPNProtocol.NORDLYNX: ProtocolFeature(
        name=VPNProtocol.NORDLYNX,
        ports={51820},
        packet_patterns=[
            {"layer": "UDP", "length": 148},
            {"layer": "UDP", "length": 92}
        ],
        tls_fingerprints=[],
        flow_characteristics={
            "min_packet_size": 92,
            "max_packet_size": 148,
            "protocol_ratio": {"UDP": 1.0},
            "packet_interval": {"min": 0.1, "max": 0.3}
        },
        detection_difficulty="高",
        evasion_techniques=["WireGuard改进版", "双NAT系统", "动态密钥协商"]
    ),
    
    VPNProtocol.TAILSCALE: ProtocolFeature(
        name=VPNProtocol.TAILSCALE,
        ports={443},
        packet_patterns=[
            {"layer": "TLS", "handshake_type": "ClientHello"},
            {"layer": "TCP", "flags": "S", "options": ["MSS", "SACK_PERM"]}
        ],
        tls_fingerprints=[
            {"cipher_suites": ["TLS_AES_128_GCM_SHA256"]},
            {"extensions": ["server_name", "status_request"]}
        ],
        flow_characteristics={
            "min_packet_size": 100,
            "max_packet_size": 1500,
            "protocol_ratio": {"TCP": 0.7, "UDP": 0.3},
            "packet_interval": {"min": 0.1, "max": 0.5}
        },
        detection_difficulty="高",
        evasion_techniques=["DERP协议", "STUN/TURN中继", "分布式节点特征"]
    )
}

def get_protocol_features(protocol: VPNProtocol) -> ProtocolFeature:
    """获取指定协议的特征"""
    return VPN_KNOWLEDGE_BASE.get(protocol)

def get_all_protocols() -> List[VPNProtocol]:
    """获取所有支持的VPN协议"""
    return list(VPN_KNOWLEDGE_BASE.keys())

def get_protocol_by_port(port: int) -> List[VPNProtocol]:
    """根据端口号获取可能的VPN协议"""
    return [
        protocol for protocol, feature in VPN_KNOWLEDGE_BASE.items()
        if port in feature.ports
    ]

def get_protocol_by_tls_fingerprint(fingerprint: Dict) -> List[VPNProtocol]:
    """根据TLS指纹获取可能的VPN协议"""
    matches = []
    for protocol, feature in VPN_KNOWLEDGE_BASE.items():
        for fp in feature.tls_fingerprints:
            if all(fp[k] == fingerprint[k] for k in fp):
                matches.append(protocol)
    return matches 